package top.ochiamalu.aioj.controller;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import top.ochiamalu.aioj.util.AIUtils;
import top.ochiamalu.aioj.util.model.AIResult;
import top.ochiamalu.aioj.util.model.AIStreamCallback;

import java.io.IOException;
import java.util.List;

/**
 * AI 对话控制器
 *
 * @author ochiamalu
 * @date 2024/03/12
 */
@Slf4j
@RestController
@RequestMapping("/ai")
public class AIController {

    /**
     * 普通对话接口
     *
     * @param message 消息
     * @return {@link AIResult }
     */
    @GetMapping("/chat/normal")
    public AIResult normalChat(@RequestParam String message) {
        return AIUtils.getInstance().sendMessage(message, false, false);
    }

    /**
     * 普通推理对话接口
     *
     * @param message 消息
     * @return {@link AIResult }
     */
    @GetMapping("/chat/normal/reasoning")
    public AIResult normalReasoningChat(@RequestParam String message) {
        return AIUtils.getInstance().sendMessage(message, false, true);
    }

    /**
     * 流式对话接口
     *
     * @param message 消息
     * @return {@link SseEmitter }
     */
    @GetMapping(value = "/chat/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamChat(@RequestParam String message) {
        return handleStreamChat(message, false);
    }

    /**
     * 流式推理对话接口
     *
     * @param message 消息
     * @return {@link SseEmitter }
     */
    @GetMapping(value = "/chat/stream/reasoning", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamReasoningChat(@RequestParam String message) {
        return handleStreamChat(message, true);
    }

    /**
     * 多轮普通对话接口
     *
     * @param messages 消息
     * @return {@link AIResult }
     */
    @PostMapping("/chat/normal/multi")
    public AIResult normalMultiChat(@RequestBody List<String> messages) {
        return AIUtils.getInstance().sendMessageList(messages, false, false);
    }

    /**
     * 多轮普通推理对话接口
     *
     * @param messages 消息
     * @return {@link AIResult }
     */
    @PostMapping("/chat/normal/multi/reasoning")
    public AIResult normalMultiReasoningChat(@RequestBody List<String> messages) {
        return AIUtils.getInstance().sendMessageList(messages, false, true);
    }

    /**
     * 多轮流式对话接口
     *
     * @param messages 消息
     * @return {@link SseEmitter }
     */
    @PostMapping(value = "/chat/stream/multi", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamMultiChat(@RequestBody List<String> messages) {
        return handleStreamMultiChat(messages, false);
    }

    /**
     * 多轮流式推理对话接口
     *
     * @param messages 消息
     * @return {@link SseEmitter }
     */
    @PostMapping(value = "/chat/stream/multi/reasoning", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamMultiReasoningChat(@RequestBody List<String> messages) {
        return handleStreamMultiChat(messages, true);
    }

    /**
     * 处理流式对话
     *
     * @param message 消息
     * @param useR1   使用r1
     * @return {@link SseEmitter }
     */
    private SseEmitter handleStreamChat(String message, boolean useR1) {
        SseEmitter emitter = new SseEmitter(5 * 60 * 1000L);
        
        emitter.onTimeout(() -> {
            log.warn("SSE连接超时");
            emitter.complete();
        });
        
        emitter.onCompletion(() -> log.info("SSE连接已完成"));

        AIUtils.getInstance().sendStreamMessage(message, useR1, new AIStreamCallback() {
            @Override
            public void onMessage(String content, String reasoningContent) {
                try {
                    if (content != null) {
                        emitter.send(SseEmitter.event()
                                .name("content")
                                .data(content));
                    }
                    if (reasoningContent != null) {
                        emitter.send(SseEmitter.event()
                                .name("reasoning")
                                .data(reasoningContent));
                    }
                } catch (IOException e) {
                    log.error("发送SSE消息时发生错误", e);
                    emitter.completeWithError(e);
                }
            }

            @Override
            public void onError(String error) {
                log.error("AI处理发生错误: {}", error);
                try {
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data(error));
                    emitter.complete();
                } catch (IOException e) {
                    log.error("发送错误消息时发生异常", e);
                    emitter.completeWithError(e);
                }
            }

            @Override
            public void onComplete() {
                log.info("AI处理完成");
                emitter.complete();
            }
        });

        return emitter;
    }

    /**
     * 处理多轮流式对话
     *
     * @param messages 消息
     * @param useR1    使用r1
     * @return {@link SseEmitter }
     */
    private SseEmitter handleStreamMultiChat(List<String> messages, boolean useR1) {
        SseEmitter emitter = new SseEmitter(5 * 60 * 1000L);
        
        emitter.onTimeout(() -> {
            log.warn("多轮对话SSE连接超时");
            emitter.complete();
        });
        
        emitter.onCompletion(() -> log.info("多轮对话SSE连接已完成"));

        AIUtils.getInstance().sendStreamMessageList(messages, useR1, new AIStreamCallback() {
            @Override
            public void onMessage(String content, String reasoningContent) {
                try {
                    if (content != null) {
                        emitter.send(SseEmitter.event()
                                .name("content")
                                .data(content));
                    }
                    if (reasoningContent != null) {
                        emitter.send(SseEmitter.event()
                                .name("reasoning")
                                .data(reasoningContent));
                    }
                } catch (IOException e) {
                    log.error("发送多轮对话SSE消息时发生错误", e);
                    emitter.completeWithError(e);
                }
            }

            @Override
            public void onError(String error) {
                log.error("多轮对话AI处理发生错误: {}", error);
                try {
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data(error));
                    emitter.complete();
                } catch (IOException e) {
                    log.error("发送错误消息时发生异常", e);
                    emitter.completeWithError(e);
                }
            }

            @Override
            public void onComplete() {
                log.info("多轮对话AI处理完成");
                emitter.complete();
            }
        });

        return emitter;
    }
} 